/*******************************************************************************
 * Copyright (c) 2014 Steffen Mikkelsen
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Public License v3.0
 * which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/gpl.html
 * 
 * Contributors:
 *     Steffen Mikkelsen - initial API and implementation
 ******************************************************************************/
package com.subsecret.secureconnection.securestream.layers.aes;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Random;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import com.subsecret.secureconnection.securestream.SecureStreamLayer;
import com.subsecret.secureconnection.securestream.layers.helper.DiffieHellmanGenerator;
import com.subsecret.tuples.Tuple;
import com.subsecret.tuples.Tuple.Tuple2;

public class AESEncryptionLayer implements SecureStreamLayer {
    
    @Override
    public int getLayerID() {
        return 20;
    }
    
    @Override
    public String getLayerDescription() {
        return "AES Encryption";
    }
    
    @Override
    public Tuple2<InputStream, OutputStream> createLayerClient(InputStream in, OutputStream out) throws IOException {
        byte[] sharedSecret = DiffieHellmanGenerator.generateKeyClient(in, out);
        Cipher aesEncryptionCipher = getAESEncryptionCipher(sharedSecret);
        Cipher aesDecryptionCipher = getAESDecryptionCipher(sharedSecret);
        in = new CipherInputStream(in, aesDecryptionCipher);
        out = new CipherOutputStream(out, aesEncryptionCipher);
        out.flush();
        return Tuple.of(in, out);
    }
    
    @Override
    public Tuple2<InputStream, OutputStream> createLayerServer(InputStream in, OutputStream out) throws IOException {
        byte[] sharedSecret = DiffieHellmanGenerator.generateKeyServer(in, out);
        Cipher aesEncryptionCipher = getAESEncryptionCipher(sharedSecret);
        Cipher aesDecryptionCipher = getAESDecryptionCipher(sharedSecret);
        out = new CipherOutputStream(out, aesEncryptionCipher);
        out.flush();
        in = new CipherInputStream(in, aesDecryptionCipher);
        return Tuple.of(in, out);
    }
    
    
    public static Cipher getAESEncryptionCipher(byte[] secret) {
        return getAESCipher(secret, Cipher.ENCRYPT_MODE);
    }
    
    public static Cipher getAESDecryptionCipher(byte[] secret) {
        return getAESCipher(secret, Cipher.DECRYPT_MODE);
    }
    
    private static Cipher getAESCipher(byte[] secret, int opmode) {
        // convert to 256 bytes to be used as AES key
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            md.update(secret);
            byte[] digest = md.digest();
            
            Key skey = new SecretKeySpec(digest, "AES");
            Cipher cipher = Cipher.getInstance("AES/CTR/NoPadding");
            int randSeed = 0;
            for (int i = 0; i < secret.length; i++) {
                randSeed += secret[i];
            }
            Random rand = new Random(randSeed);
            byte[] IV = new byte[16];
            rand.nextBytes(IV);
            cipher.init(opmode, skey, new IvParameterSpec(IV));
            return cipher;
        }
        catch (NoSuchAlgorithmException | InvalidKeyException | InvalidAlgorithmParameterException | NoSuchPaddingException e) {
            throw new RuntimeException(e.getClass().getSimpleName() + ": " + e.getMessage());  
        }
    }
    
}
